package com.runmyai.aitools.service;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.github.mustachejava.DefaultMustacheFactory;
import com.github.mustachejava.Mustache;
import com.github.mustachejava.MustacheFactory;
import com.runmyai.aitools.exception.ResponseErrorException;
import com.runmyai.aitools.exception.TokenErrorException;
import com.runmyai.aitools.model.AiRequest;
import com.runmyai.aitools.model.AiResponse;
import com.runmyai.aitools.model.ContentInfo;
import com.runmyai.aitools.text.TextParsingDispatcher;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.stereotype.Service;
import org.springframework.util.Assert;

import java.io.IOException;
import java.io.StringReader;
import java.io.StringWriter;
import java.io.Writer;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @author wry
 */
@Slf4j
@Service
public class AiService {

    /**
     * 百度云千帆大模型开发平台API_KEY和SECRET_KEY
     */
    private static final String API_KEY = "";
    private static final String SECRET_KEY = "";

    /**
     * ERNIE-Bot 4.0 API
     */
    private static final String ERNIE_API = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=";


    private final OkHttpClient client;
    private static final MediaType MEDIA_TYPE = MediaType.parse("application/json");
    private static final RequestBody TOKEN_REQUEST_BODY = RequestBody.create(MEDIA_TYPE, "{}");

    private static final String ERROR_CODE_KEY = "error_code";
    private static final String ERROR_MSG_KEY = "error_msg";
    private static final String RESULT_KEY = "result";
    private static final String NORMAL_KEY = "normal";

    private static final int MAX_TEXT_LENGTH = 4500;
    private static final int MAX_RETRIES = 3;

    private String tokenUrl;
    private String token;

    public AiService() {
        this.client = new OkHttpClient.Builder()
                .connectTimeout(10, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .readTimeout(90, TimeUnit.SECONDS)
                .build();
        this.tokenUrl = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=" + API_KEY + "&client_secret=" + SECRET_KEY;
        try {
            this.token = getAccessToken();
        } catch (TokenErrorException e) {
            log.error("获取ai token失败: {}", e.getMessage());
        }
    }


    /**
     * 发起新的AI会话
     *
     * @param promptTmpl 提示词模板
     * @param content    请求参数
     * @param num        生成数量
     * @return
     */
    public AiResponse startAigc(String promptTmpl, ContentInfo content, Integer num) {
        AiRequest aiRequest = buildRequestBody(promptTmpl, content, num);
        return processAiRequest(aiRequest);
    }

    /**
     * String prompt = "继续生成";
     * 继续进行会话, 带有上下文.
     *
     * @param aiRequest AI会话的上下文
     * @param prompt    用户提问的问题
     * @return
     * @throws RuntimeException
     */
    public AiResponse continueAigc(AiRequest aiRequest, String prompt) {
        aiRequest.query(prompt);
        return processAiRequest(aiRequest);
    }


    /**
     * 调用AI服务并解析返回的数据, 互记录 <Chat> 会保存在AiResponse
     *
     * @param aiRequest AI请求对象
     * @return aiResponse AI响应对象
     */
    public AiResponse processAiRequest(AiRequest aiRequest) {
        List<String> errors = null;
        String responseText = null;
        List<String> data = null;

        // 发送AI请求并获得响应
        Response response = null;
        try {
            response = sendAiRequest(aiRequest);
            // 处理获得的响应，并获取响应文本. 进程可能于当前阶段中止
            JSONObject respJson = handleResponse(response, aiRequest, 0);

            // 处理返回结果
            responseText = respJson.getString(RESULT_KEY);

            String finishReason = respJson.getString("finish_reason");
            errors = buildErrorMsg(respJson, finishReason);

            if (!NORMAL_KEY.equals(finishReason)) {
                return new AiResponse(aiRequest.getMessages(), responseText, null, errors);
            }

            // 使用文本解析器解析响应文本，获取数据列表
            data = TextParsingDispatcher.parse(responseText);

            if (data != null) {
                data = data.stream().distinct().collect(Collectors.toList());
            }
        } catch (ResponseErrorException | TokenErrorException | IOException error) {
            log.error("处理AI任务失败: ", error);
            if (errors == null) {
                errors = Arrays.asList(error.getMessage());
            } else {
                errors.add(error.getMessage());
            }
        } finally {
            if (response != null) {
                response.close(); // 如果Response实现了AutoCloseable，这将自动关闭它
            }
        }

        // 使用请求的消息、响应文本和解析后的数据创建并返回一个新的AI响应对象
        return new AiResponse(aiRequest.getMessages(), responseText, data, errors);
    }

    private List<String> buildErrorMsg(JSONObject respJson, String finishReason) {
        List<String> errors = new ArrayList<>();
        boolean isTruncated = respJson.getBooleanValue("is_truncated");
        boolean needClearHistory = respJson.getBooleanValue("need_clear_history");

        if (needClearHistory) {
            int banRound = respJson.getIntValue("ban_round");
            if (banRound == -1) {
                errors.add("用户输入内容有敏感信息, 建议关闭当前会话");
            } else {
                errors.add("第" + banRound + "轮对话有敏感信息, 建议关闭当前会话并清理会话信息");
            }
            return errors;
        }
        if (isTruncated) {
            errors.add("输出结果超过长度限制");
        } else if ("stop".equals(finishReason)) {
            errors.add("输出结果命中入参stop中指定的字段后被截断");
        } else if ("content_filter".equals(finishReason)) {
            errors.add("输出内容被截断、兜底、替换为**等");
        }
        return errors;
    }

    /**
     * 负责发起AI请求并得到Response。
     *
     * @param aiRequest AI请求对象
     * @return response Response
     * @throws TokenErrorException 如果发生TOKEN错误
     * @throws IOException         如果发生I/O错误
     */
    private Response sendAiRequest(AiRequest aiRequest) throws TokenErrorException, IOException {
        // 确保aiRequest中的messages数量是奇数
        Assert.state(aiRequest.getMessages().size() % 2 == 1, "请求数据异常! messages必须是奇数个");

        // 创建请求体，将aiRequest对象转换为JSON字符串
        RequestBody body = RequestBody.create(MEDIA_TYPE, JSON.toJSONString(aiRequest));

        // 构建请求对象
        Request request = buildRequest(body);

        try {
            Response response = client.newCall(request).execute();
            // 执行请求并返回响应
            return response;
        } catch (Exception e) {
            throw e;
        }
    }

    /**
     * 根据给定的提示和内容构建AI请求对象。
     *
     * @param promptTmpl 提示语模板
     * @param content    内容信息
     * @return 构建好的AI请求对象
     */
    private AiRequest buildRequestBody(String promptTmpl, ContentInfo content, Integer num) {
        // 清理content中的不必要数据或格式
        content.clean();

        // 根据需要修剪content的内容
        trimContentIfNeeded(content, promptTmpl);

        // 解析提示模板并获取实际的提示内容
        String prompt = parseTemplate(promptTmpl, buildTemplateParamObject(content, num));

        log.info("提示词内容: {}", prompt);

        return new AiRequest(prompt);
    }

    private Object buildTemplateParamObject(ContentInfo content, Integer num) {
        Map<String, String> obj = new HashMap<>(8);
        obj.put("content", content.getContent());
        if (num != null) {
            Assert.state(num > 0, "参数异常! 数量必须大于0");
            obj.put("num", num.toString());
        } else {
            obj.put("num", "10");
        }
        return obj;
    }

    /**
     * 解析提示词模板
     *
     * @param template
     * @param param
     * @return
     */
    public static String parseTemplate(String template, Object param) {
        MustacheFactory mf = new DefaultMustacheFactory();
        Mustache mustache = mf.compile(new StringReader(template), "");
        try (Writer writer = new StringWriter()) {
            mustache.execute(writer, param).flush();
            return writer.toString();
        } catch (IOException e) {
            throw new IllegalArgumentException("模板解析异常");
        }
    }


    /**
     * 根据需要对内容进行修剪
     */
    public void trimContentIfNeeded(ContentInfo content, @NonNull String prompt) {
        // 计算提示词的大概总长度
        int totalLength = content.getContentLength() + prompt.length();

        // 如果总长度超过4500个字符, 文心一言要求长度小于4800且token数小于2000
        if (totalLength > MAX_TEXT_LENGTH) {

            int charsToTrim = totalLength - MAX_TEXT_LENGTH;

            // 如果需要修剪的字符数量小于或等于内容的长度
            if (charsToTrim <= content.getContentLength()) {
                // 对内容进行修剪，去掉末尾的部分字符
                content.setContent(content.getContent().substring(0, content.getContent().length() - charsToTrim));
            } else {
                // 如果内容的长度小于需要截掉的字符数量，则将内容设置为空字符串
                content.setContent("");
            }
        }
    }

    /**
     * 创建 request
     *
     * @param body 请求body
     * @return
     * @throws TokenErrorException
     */
    private Request buildRequest(RequestBody body) throws TokenErrorException {
        if (this.token == null) {
            this.token = getAccessToken();
        }
        return new Request.Builder()
                .url(ERNIE_API + this.token)
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .build();
    }

    /**
     * 对AI接口返回的response进行解析。发现token过期，则执行重新获取response并进行解析的流程。
     *
     * @param response   接口返回对象
     * @param aiRequest  请求数据
     * @param retryCount 重试次数
     * @return
     * @throws IOException
     * @throws TokenErrorException
     * @throws ResponseErrorException
     */
    private JSONObject handleResponse(Response response, AiRequest aiRequest, int retryCount) throws IOException, TokenErrorException, ResponseErrorException {
        if (!response.isSuccessful()) {
            throw new ResponseErrorException(response.message());
        }

        ResponseBody respBody = response.body();

        String string = respBody.string();

        JSONObject jsonObject = JSON.parseObject(string);

        if (jsonObject.containsKey(ERROR_CODE_KEY)) {
            int errorCode = jsonObject.getInteger(ERROR_CODE_KEY);
            if (errorCode == 110) {
                // 请求token, 然后重新发起请求
                return refreshTokenAndRetry(aiRequest, retryCount);
            }
            handleError(jsonObject);
        }
        return jsonObject;
    }

    private void handleError(JSONObject jsonObject) throws ResponseErrorException {
        int errorCode = jsonObject.getInteger(ERROR_CODE_KEY);
        String errorMsg = jsonObject.getString(ERROR_MSG_KEY);
        log.error("AI Response error: {}, message: {}", errorCode, errorMsg);
        throw new ResponseErrorException(errorMsg);
    }

    /**
     * 重新获取response并进行解析
     *
     * @param aiRequest  AI请求参数
     * @param retryCount 重试次数
     * @return
     * @throws TokenErrorException
     * @throws ResponseErrorException
     * @throws IOException
     */
    private JSONObject refreshTokenAndRetry(AiRequest aiRequest, int retryCount) throws TokenErrorException, ResponseErrorException, IOException {
        this.token = null;
        retryCount = retryCount + 1;
        if (retryCount > MAX_RETRIES) {
            throw new TokenErrorException("Token过期, 已超过重试次数");
        }
        Response response = sendAiRequest(aiRequest);
        return handleResponse(response, aiRequest, retryCount);
    }

    private String getAccessToken() throws TokenErrorException {
        Request request = new Request.Builder()
                .url(tokenUrl)
                .post(TOKEN_REQUEST_BODY)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .build();

        try (Response response = this.client.newCall(request).execute()) {
            JSONObject jsonObject = JSON.parseObject(response.body().string());
            return jsonObject.getString("access_token");
        } catch (IOException e) {
            log.error("获取token失败: ", e.getCause());
            throw new TokenErrorException("Token 获取失败! ", e.getCause());
        }
    }

/*
    public static void main(String[] args) {
        AiService aiService = new AiService();
        ContentInfo contentInfo = new ContentInfo("来自沙县博物馆的藏品");
        AiResponse aiResponse = aiService.startAigc("{{content}}，请根据上述提示词，生成 {{num}}条内容", contentInfo, 5);
        List<String> data = aiResponse.getData();
        data.forEach(item -> System.out.println(item));
    }
*/

}
